import cv2
import numpy as np
from matplotlib import pyplot as plt
import os
import sys
from PIL import Image

# face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt_tree.xml')
# face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalcatface.xml')
face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt.xml')

recognizer = cv2.face.LBPHFaceRecognizer.create()


def getImagesAndLabels(base_path):
    imagePaths = [os.path.join(base_path, f) for f in os.listdir(base_path)]
    faceSamples = []
    ids = []
    for imagePath in imagePaths:
        PIL_img = Image.open(imagePath).convert('L')
        img_numpy = np.array(PIL_img, 'uint8')
        id = int(os.path.split(imagePath)[-1].split('.')[0])
        faces = face_detector.detectMultiScale(img_numpy)

        for x, y, w, h in faces:
            faceSamples.append(img_numpy[y: y+h, x: x+w])
            ids.append(id)

    return faceSamples, ids


if __name__ == '__main__':
    path = './data/03/'
    faces, ids = getImagesAndLabels(path)
    recognizer.train(faces, np.array(ids))
    # 保存模型到 yml 文件
    recognizer.write('./data/03.yml')
    print('完成')
